import { MilvusClient, DataType } from '@zilliz/milvus2-sdk-node'
import { config } from '../config/index.js'
import { modelService } from './model.js'  // 导入 modelService
import logger from './logger.js'

export class VectorStore {
  constructor() {
    this.maxRetries = 3
    this.retryDelay = 1000
    this.database = 'default'
    this.collectionName = 'documents'
    this.dimension = 768 // nomic-embed-text 的向量维度是 768
    this.isInitialized = false
  }

  initClient() {
    this.client = new MilvusClient({
      address: config.milvus.address || 'localhost:19530',
      username: config.milvus.username,
      password: config.milvus.password,
      ssl: false,
      connectTimeout: 30000,  // 增加超时时间
      receiveTimeout: 30000,
      keepAlive: true,
      keepAliveTimeout: 30000
    })
  }

  async init() {
    if (this.isInitialized) {
      return true
    }

    try {
      logger.info('Initializing vector store...')
      
      // 1. 初始化客户端
      this.initClient()
      
      // 2. 等待 Milvus 服务就绪
      await this.waitForMilvus()

      // 3. 创建或加载集合
      await this.ensureCollection()

      this.isInitialized = true
      return true
    } catch (error) {
      logger.error('Failed to initialize vector store:', {
        error: error.message,
        type: error.constructor.name,
        collection: this.collectionName
      })
      throw error
    }
  }

  async waitForMilvus(maxAttempts = 10) {
    for (let i = 0; i < maxAttempts; i++) {
      try {
        await this.client.listCollections()
        logger.info('Milvus service is ready')
        return
      } catch (error) {
        logger.warn(`Waiting for Milvus service (attempt ${i + 1}/${maxAttempts})`)
        await new Promise(resolve => setTimeout(resolve, 3000))
      }
    }
    throw new Error('Milvus service not available')
  }

  async ensureCollection() {
    try {
      // 检查集合是否存在
      const collections = await this.client.listCollections()
      const exists = collections.collection_names?.includes(this.collectionName)

      if (!exists) {
        logger.info('Creating new collection:', this.collectionName)
        await this.createCollection()
        await new Promise(resolve => setTimeout(resolve, 3000))
        
        logger.info('Creating index...')
        await this.createIndex()
        await new Promise(resolve => setTimeout(resolve, 3000))
      }

      // 加载集合
      await this.client.loadCollection({
        collection_name: this.collectionName
      })

      // 验证集合状态
      const stats = await this.client.getCollectionStatistics({
        collection_name: this.collectionName
      })

      const desc = await this.client.describeCollection({
        collection_name: this.collectionName
      })

      logger.info('Collection ready:', {
        name: this.collectionName,
        rowCount: stats.row_count,
        dimension: this.dimension,
        hasIndex: desc.indexes?.length > 0
      })

    } catch (error) {
      logger.error('Failed to ensure collection:', error)
      throw error
    }
  }

  async createCollection() {
    try {
      // 先尝试删除可能存在的同名集合
      try {
        await this.client.dropCollection({
          collection_name: this.collectionName
        })
        logger.info('Dropped existing collection')
      } catch (error) {
        // 忽略删除错误
      }

      // 创建新集合
      await this.client.createCollection({
        collection_name: this.collectionName,
        dimension: this.dimension,
        fields: [
          {
            name: 'id',
            description: 'Chunk ID',
            data_type: DataType.VarChar,
            is_primary_key: true,
            max_length: 100
          },
          {
            name: 'parent_id',
            description: 'Parent Document ID',
            data_type: DataType.VarChar,
            max_length: 100
          },
          {
            name: 'content',
            description: 'Chunk Content',
            data_type: DataType.VarChar,
            max_length: 65535
          },
          {
            name: 'title',
            description: 'Document Title',
            data_type: DataType.VarChar,
            max_length: 255
          },
          {
            name: 'chunk_index',
            description: 'Chunk Index',
            data_type: DataType.Int64
          },
          {
            name: 'embedding',
            description: 'Text Embedding',
            data_type: DataType.FloatVector,
            dim: this.dimension
          }
        ]
      })

      logger.info('Collection created successfully')
    } catch (error) {
      logger.error('Failed to create collection:', {
        error: error.message,
        collection: this.collectionName
      })
      throw error
    }
  }

  async createIndex() {
    try {
      logger.info('Creating index...')
      
      await this.client.createIndex({
        collection_name: this.collectionName,
        field_name: 'embedding',
        index_type: 'IVF_FLAT',
        metric_type: 'L2',
        params: { nlist: 2048 }
      })

      logger.info('Index created successfully')
    } catch (error) {
      logger.error('Failed to create index:', error)
      throw error
    }
  }

  async insert(chunks) {
    return this.withRetry(() => this._insert(chunks))
  }

  async _insert(chunks) {
    try {
      const chunksArray = Array.isArray(chunks) ? chunks : [chunks]
      
      logger.info('Processing chunks for insertion:', {
        count: chunksArray.length,
        firstChunk: {
          id: chunksArray[0]?.id,
          title: chunksArray[0]?.title,
          contentLength: chunksArray[0]?.content?.length,
          embeddingLength: chunksArray[0]?.embedding?.length
        }
      })

      // 验证向量维度
      const firstEmbedding = chunksArray[0]?.embedding
      if (!firstEmbedding || firstEmbedding.length !== this.dimension) {
        throw new Error(`Invalid embedding dimension: expected ${this.dimension}, got ${firstEmbedding?.length}`)
      }

      // 格式化数据
      const formattedChunks = chunksArray.map(chunk => {
        // 验证必要字段
        if (!chunk.id || !chunk.content || !chunk.embedding) {
          throw new Error('Missing required fields in chunk')
        }

        // 验证并格式化向量
        const embedding = chunk.embedding.map(num => {
          const n = Number(num)
          if (isNaN(n) || !isFinite(n)) {
            throw new Error('Invalid embedding value')
          }
          return Number(n.toFixed(6))
        })

        return {
          id: String(chunk.id),
          parent_id: String(chunk.parent_id || ''),
          title: String(chunk.title || ''),
          content: String(chunk.content),
          embedding: embedding,
          chunk_index: Number(chunk.id.split('_')[1] || 0)
        }
      })

      const insertData = {
        collection_name: this.collectionName,
        fields_data: formattedChunks
      }

      // 执行插入
      const result = await this.client.insert(insertData)

      // 验证插入结果
      logger.info('Insert result:', {
        success: result.succ_count,
        expected: chunksArray.length,
        insertIds: result.ids
      })

      // 刷新集合以确保数据可见
      await this.client.flush({
        collection_names: [this.collectionName]
      })

      // 验证数据是否可查询
      const stats = await this.client.getCollectionStatistics({
        collection_name: this.collectionName
      })

      logger.info('Collection stats after insert:', {
        rowCount: stats.row_count
      })

      return result
    } catch (error) {
      logger.error('Insert failed:', {
        error: error.message,
        type: error.constructor.name,
        stack: error.stack?.split('\n')[0]
      })
      throw error
    }
  }

  async search(embedding, options = {}) {
    try {
      // 1. 验证输入
      if (!Array.isArray(embedding)) {
        throw new Error('Embedding must be an array')
      }

      if (embedding.length !== this.dimension) {
        throw new Error(`Invalid embedding dimension: expected ${this.dimension}, got ${embedding.length}`)
      }

      // 2. 确保向量数据是数字数组
      const normalizedEmbedding = embedding.map(n => {
        const num = Number(n)
        if (isNaN(num)) {
          throw new Error('Embedding contains non-numeric values')
        }
        return Number(num.toFixed(6)) // 限制精度
      })

      // 3. 加载集合
      await this.client.loadCollection({
        collection_name: this.collectionName
      })

      // 4. 准备搜索参数
      const searchParams = {
        collection_name: this.collectionName,
        vector: normalizedEmbedding,
        output_fields: ["id", "title", "content", "parent_id", "chunk_index"],
        top_k: config.ollama.parameters.top_k,
        limit: options.limit || 3,
        metric_type: "L2",        // L2 距离越小表示越相似
      }

      // 5. 执行搜索
      const response = await this.client.search(searchParams)

      // 6. 验证结果
      if (!response || !response.results) {
        logger.warn('Empty search response')
        return { data: [], total: 0 }
      }

      // 7. 格式化结果
      // Milvus 已经按 L2 距离升序排序了结果（距离小的在前）
      const results = response.results.map(result => ({
        id: result.id || '',
        title: result.title || '',
        content: result.content || '',
        parent_id: result.parent_id || '',
        chunk_index: result.chunk_index || 0,
        score: 100 / (100 + (result.score || 0)) // 将 L2 距离转换为相似度分数（用于显示）
      }))

      return {
        data: results,
        total: results.length
      }

    } catch (error) {
      logger.error('Search failed:', {
        error: error.message,
        type: error.constructor.name,
        collection: this.collectionName,
        stack: error.stack
      })
      throw error
    }
  }

  async deleteVectors(expr) {
    try {
      logger.info('Deleting vectors with filter:', expr)
      
      await this.client.loadCollection({
        collection_name: this.collectionName
      })

      const result = await this.client.delete({
        collection_name: this.collectionName,
        filter: expr
      })

      logger.info('Delete result:', result)
      return result
    } catch (error) {
      logger.error('Failed to delete from vector store:', {
        filter: expr,
        error: error.message,
        stack: error.stack
      })
      throw error
    }
  }

  async query(options = {}) {
    try {
      await this.client.loadCollection({
        collection_name: this.collectionName
      })

      const result = await this.client.query({
        collection_name: this.collectionName,
        output_fields: options.output_fields || ['id', 'title', 'content'],
        limit: options.limit || 100
      })

      logger.info('Query result:', {
        resultLength: result.data?.length || 0,
        firstDoc: result.data?.[0]
      })

      return result.data || []
    } catch (error) {
      logger.error('Failed to query:', error)
      throw error
    }
  }

  async getDocumentChunks(documentId) {
    try {
      const response = await this.client.query({
        collection_name: this.collectionName,
        expr: `parent_id == "${documentId}"`,
        output_fields: ["id", "title", "content", "chunk_index"],
        consistency_level: "Strong"
      })

      if (!response || !response.data) {
        return []
      }

      // 按 chunk_index 排序
      return response.data.sort((a, b) => a.chunk_index - b.chunk_index)
    } catch (error) {
      logger.error('Failed to get document chunks:', {
        documentId,
        error: error.message
      })
      throw error
    }
  }

  async withRetry(operation) {
    let retries = this.maxRetries
    while (retries > 0) {
      try {
        return await operation()
      } catch (error) {
        retries--
        if (retries === 0) throw error
        logger.warn('Operation failed, retrying...', { error: error.message, retriesLeft: retries })
        await new Promise(resolve => setTimeout(resolve, this.retryDelay))
        await this.reconnect()
      }
    }
  }

  async reconnect() {
    try {
      await this.client.closeConnection()
    } catch (err) {
      logger.warn('Error closing connection:', err)
    }
    this.initClient()
    await this.init()
  }

  async insertVectors(vectors) {
    try {
      await this.init()
      
      // 验证向量数据
      const validVectors = vectors.filter(vector => {
        if (!vector.embedding || vector.embedding.length !== this.dimension) {
          logger.warn('Invalid vector data:', {
            id: vector.id,
            embeddingLength: vector.embedding?.length
          })
          return false
        }
        return true
      })
      
      if (validVectors.length === 0) {
        logger.warn('No valid vectors to insert')
        return
      }
      
      // 插入向量
      await this.client.insert({
        collection_name: this.collectionName,
        data: validVectors
      })
      
      // 执行 flush
      await this.client.flush({
        collection_names: [this.collectionName]
      })
      
      logger.info('Vectors inserted successfully:', {
        count: validVectors.length
      })
      
      return validVectors.length
    } catch (error) {
      logger.error('Failed to insert vectors:', error)
      throw error
    }
  }

}

// 修改初始化方式
export const vectorStore = new VectorStore()